import os
import sys
import numpy as np

from PIL import Image
from pathlib import Path


if __name__ == "__main__":


    STEP_VIDEO_CONVERT_25FPS = False
    STEP_VIDEO_EXTRACT_FRAME = False
    STEP_VIDEO_EXTRACT_3DMM  = False

    STEP_AUDIO_TO_EXPRESSION = False

    STEP_RENDER_BLEND = False
    STEP_RENDER_GENERATE_LIST = False
    STEP_RENDER_ARCFACE = False
    STEP_RENDER_TRAIN = True
    STEP_RENDER_TEST = False

    path_data = Path("../data")
    path_checkpoints = Path("../checkpoints")
    path_expr = Path("../experiments")
    path_work = Path("/mnt/data/CODE/Audio-driven-TalkingFace-HeadPose/reconstruction")


    ########################################### 
    ## prepare data
    path_audio = path_data.joinpath('input/audio/5_00006.wav')
    path_video = path_data.joinpath('input/video/31.mp4')

    video_name = path_video.stem
    audio_name = path_audio.stem

    path_save = path_data.joinpath(f"processing/{video_name}")
    path_audio_save = path_save.joinpath(f'audio')
    path_video_save = path_save.joinpath(f'video')
    path_video_fps = path_video_save.joinpath(f'video.mp4')

    path_audio_save.mkdir(parents=True, exist_ok=True)
    path_video_save.mkdir(parents=True, exist_ok=True)


    ### convert the target video to 25fps
    if STEP_VIDEO_CONVERT_25FPS:
        cmd = f"cd video; "\
              f"ffmpeg -i {path_video} -r 25 {path_video_fps}; "
        print(f"[xxxx] RUN: \n [xxxx] {cmd}")
        os.system(cmd)

        # cmd = f"cd video; "\
        #       f"ffmpeg -i tmp.mp4 -vf scale=\"720:-1\" {path_video_fps}"
        # print(f"[xxxx] RUN: \n [xxxx] {cmd}")
        # os.system(cmd)

    ### extract frames and landmarks by
    if STEP_VIDEO_EXTRACT_FRAME:
        cmd = f"cd utils; "\
              f"python extract_frames.py --path_video {path_video_fps} --path_save {path_save}"

        print(f"[xxxx] RUN: \n [xxxx] {cmd}")
        os.system(cmd)

        # extract landarks
        cmd = f"cd utils; "\
              f"python extract_landmarks.py --path_save {path_save} "
        print(f"[xxxx] RUN: \n [xxxx] {cmd}")
        os.system(cmd)


    ### extract 3DMM coeffs per frames
    if STEP_VIDEO_EXTRACT_3DMM:
        cmd = f"cd reconstruction;"\
              f"CUDA_VISIBLE_DEVICES=0 python reconstruct.py --path_data {path_save} --path_work {path_work}"

        print(f"[xxxx] RUN: \n [xxxx] {cmd}")
        os.system(cmd)

    ###########################################
    # STEP 1： Audio-to-Expression
    if STEP_AUDIO_TO_EXPRESSION:
        cmd = f"cd audio;" \
            f"python src/infer_atcnet.py --path_data {path_audio} --path_save {path_audio_save}"

        print(f"[xxxxx] RUN: \n[xxxxx] {cmd}")
        os.system(cmd)

    ############################################
    # STEP 3:  Render Network train
    if STEP_RENDER_BLEND:
        cmd = f"cd render; "\
              f"python scripts/S1_blend_images.py --path_save  {path_save}"

        print(f"[xxxxx] RUN: \n[xxxxx] {cmd}")
        os.system(cmd)

    if STEP_RENDER_GENERATE_LIST:
        cmd = f"cd render; "\
              f"python scripts/S2_get_image_list.py --path_save  {path_save}"

        print(f"[xxxxx] RUN: \n[xxxxx] {cmd}")
        os.system(cmd)

    if STEP_RENDER_ARCFACE:
        cmd = f"cd arcface; "\
              f"python inference.py --list_name trainB --path_save  {path_save}"

        print(f"[xxxxx] RUN: \n[xxxxx] {cmd}")
        os.system(cmd)

        cmd = f"cd arcface; "\
              f"python inference.py --list_name testB --path_save  {path_save}"

        print(f"[xxxxx] RUN: \n[xxxxx] {cmd}")
        os.system(cmd)

    if STEP_RENDER_TRAIN:
        cmd = f"cd render;                                        \
                python train.py                          \
                       --path_data {path_save}              \
                       --checkpoints_dir {path_checkpoints}       \
                       --name memory_seq_p2p/{video_name}         \
                       --model memory_seq          \
                       --epoch 0                   \
                       --batch_size  8             \
                       --epoch_begin 1             \
                       --lambda_mask 2             \
                       --lr 0.0001                 \
                       --display_env memory_seq_31 \
                       --gpu_ids 0                 \
                       --niter 60                  \
                       --niter_decay 0 "

        print(f"[xxxxx] RUN: \n[xxxxx] {cmd}")
        os.system(cmd)


    epoch = 60
    if STEP_RENDER_TEST:
        cmd = f"cd render;                          \
                python test.py                      \
                    --path_video {path_video_fps}              \
                    --path_data {path_video_save}              \
                    --checkpoints_dir {path_checkpoints}       \
                    --name memory_seq_p2p/{video_name}        \
                    --model memory_seq              \
                    --num_test 200                  \
                    --epoch {epoch}                \
                    --gpu_ids 0                     \
                    --imagefolder images{epoch}"
        print(f"[xxxxx] RUN: \n[xxxxx] {cmd}")
        os.system(cmd)
